"""
Utility functions for detrital modeling.
"""


__author__="Boris Avdeev, borisaqua@gmail.com"



import pymc as pm
import numpy as np
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt

def sig2tau(sig):
    """Standard deviation to precision"""
    return 1./(sig**2)

def tau2sig(tau):
    """Precision to standard deviation"""
    return (1./tau)**.5

def sample_wr(population, size):
    """Chooses 'size' random elements (with replacement) from a population"""
    n = len(population)
    js = np.array(np.random.random(size) * n).astype('int') 
    return population[js]


def parm2ha(parm, ub, ub_type):
    """Convert parameter vector of an exhumation model [ e1, e2, e3, ... ,hc, abr1, abr2, ... ]
    into ([elevation], [age]) array for plotting"""
    parm = np.array(parm).flatten()
    segs = len(parm)/2
    e = parm[0:segs]
    hc= parm[segs]
    a = [0]
    h = [hc]
    if segs > 1:
        a.extend(parm[segs+1:])
        da = np.diff(a)
        for i in range(segs-1):
            h.append(h[i] + e[i]*da[i])
    if not (ub_type=='h' or ub_type=='a'): 
        raise NameError('ub_type must be either a or h')
    if ub_type=='h' and ub > h[-1]:
        h.append(ub) 
        a.append(a[-1] + (ub - h[-2]) / e[-1])
    elif ub_type=='a' and ub > a[-1]:
        a.append(ub) 
        h.append(h[-1] + (ub - a[-2]) * e[-1])
    return np.array(h), np.array(a)


def h2a(hyps, parm):
    """Convert elevations to ages given parm = [ e1, e2, e3, ... ,hc, abr1, abr2, ... ] """
    h,a = parm2ha(parm,max(hyps),'h')
    if len(a)==1: return np.repeat(np.nan,len(hyps)) 
    fun = interp1d(h, a, bounds_error=False, fill_value=np.nan) #or 0???
    return fun(hyps)

def a2h(ages, parm):
    """Convert ages to elevations given parm = [ e1, e2, e3, ... ,hc, abr1, abr2, ... ] """
    [h,a] = parm2ha(parm,max(ages),'a')
    fun = interp1d(a, h, bounds_error=False, fill_value=0)
    return fun(ages)
    
@pm.randomwrap
def detrital_random(hyps,hyps_w,parm,err,size=None):
    """Random detrital sample. hyps -- bins. BAD? Should probably
       supply weighting function and DEM and sample from there."""
    if len(hyps_w)==1: idx=range(len(hyps))  #WARNING!
    else: idx = hyps_w > pm.runiform(0,max(hyps_w),len(hyps_w))
    ta = h2a(hyps[idx],parm)
    ta_samp = sample_wr(ta, size)    
    return pm.rnormal(ta_samp,sig2tau(err*ta_samp))





#-----DATA INPUT-------------------------------------------

def bins(hyps):
    h_w, h_edges = np.histogram(hyps, bins=512, normed=False)
    h_w = h_w / float(len(hyps))
    h_b  = h_edges[:-1] + np.diff(h_edges)/2  # Replace bin margins with bin centers
    return {'h':h_b, 'w':h_w, 'hyps': hyps}

def read_xyhs(xyhs,delimiter=',',skiprows=1):
    data  = np.loadtxt(xyhs, delimiter=delimiter, skip_header=skiprows)
    hyps  = data[:,2]
    if max(hyps) > 10:
      print "WARNING: converting to km"
      hyps /= 1000.0
    slope = data[:,3]
    hs_w, h_edges, s_edges = np.histogram2d(hyps, slope, bins=[256,64], normed=False)
    hs_w = hs_w / float(len(hyps))
    h_b  = h_edges[:-1] + np.diff(h_edges)/2  # Replace bin margins with bin centers
    s_b = s_edges[:-1] + np.diff(s_edges)/2  # Replace bin margins with bin centers
    return {'h':h_b, 's':s_b, 'w':hs_w, 'hyps': hyps}

def read_xyh(xyhs,delimiter=',',skiprows=1):
    data  = np.loadtxt(xyhs, delimiter=delimiter, skiprows=skiprows)
    hyps  = data[:,2]
    if max(hyps) > 10:
      print "WARNING: converting to km"
      hyps /= 1000.0
    return bins(hyps)

def read_h(h_file):
    hyps  = np.loadtxt(h_file)
    if max(hyps) > 10:
      print "WARNING: converting to km"
      hyps /= 1000.0
    return bins(hyps)




#-----PLOTTING FUNCTIONS---------------------------------------

def lines(observations, a=None, b=None, color='black'):
    def __ecdf(x):
        counter = 0.0
        for obs in observations:
            if obs <= x:
                counter += 1
        return counter / len(observations)

    if a == None:
        a = observations.min() - observations.std()/2
    if b == None:
        b = observations.max() + observations.std()/2
    X = np.linspace(a, b, 100)
    f = np.vectorize(__ecdf)
    plt.plot(X, f(X), color=color,alpha=0.01)
    
def points(obs,hpd=None,color='black'):
    obs=np.array(obs)
    idx = np.argsort(obs)
    cdf = np.linspace(0,1,len(obs))
    if hpd is None:
        plt.plot(obs[idx],cdf,'.',color=color)
    else:
        plt.errorbar(obs[idx],cdf,xerr=np.abs(hpd[:,idx]-obs[idx]),fmt='.')


def gof(obs,sim,obs_hpd=None, obs_color='black',sim_color='0.3', color=None):
    if color is not None:
        obs_color, sim_color = color, color
    for sim_i in sim:
        lines(sim_i,color=sim_color)
    points(obs,obs_hpd,color=obs_color)
    
 
def ah(parms, h_max, color='black'):
    for parm in parms:
        h, a = parm2ha(parm, h_max, 'h')
        plt.plot(a,h,color=color,alpha=0.01)



   

